package cn.ultrarisk.booyo.spider.http;

import cn.ultrarisk.booyo.core.spider.Proxy;
import cn.ultrarisk.booyo.core.util.CollectionUtil;
import org.apache.http.HttpHeaders;
import org.apache.http.HttpHost;
import org.apache.http.NameValuePair;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.entity.UrlEncodedFormEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.message.BasicNameValuePair;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Http下载器
 *
 * @author sevendlong(lin23871@163.com)
 * @date 2014/9/18
 * @since V1.0
 */
public class Downloader {
    /**
     * 下载页面内容
     *
     * @param url     目标地址
     * @param request 请求对象
     * @return 返回响应对象
     * @throws IOException
     */
    public static Response download(String url, Request request) throws IOException {
        CloseableHttpClient client = null;
        CloseableHttpResponse response = null;
        try {
            HttpRequestBase httpRequest;
            if (Request.Method.GET.getCode().equals(request.getMethod().getCode())) {
                httpRequest = new HttpGet(url + (url.contains("?") ? "&" : "?") + getUrlParam(request.getParams()));
            } else if (Request.Method.POST.getCode().equals(request.getMethod().getCode())) {
                httpRequest = new HttpPost(url);
                ((HttpPost) httpRequest).setEntity(getFormParam(request.getParams()));
            } else {
                return null;
            }

            httpRequest.setConfig(buildRequestConfig(request.getProxy(), request.getTimeout()));
            addHeaders(httpRequest, request.getHeaders());
            addCookies(httpRequest, request.getCookies());
            addContentType(httpRequest, request.getContentType());

            client = buildHttpClient(request.getProxy());
            response = client.execute(httpRequest);
            return new Response(response);
        } finally {
            if (client != null) {
                client.close();
            }
            if (response != null) {
                response.close();
            }
        }
    }

    /**
     * 获取get请求的url参数
     *
     * @param params 请求参数
     * @return
     * @throws UnsupportedEncodingException
     */
    private static String getUrlParam(Map<String, String[]> params) throws UnsupportedEncodingException {
        if (params == null || params.size() == 0) {
            return "";
        }

        List<String> list = new ArrayList<String>();
        for (String key : params.keySet()) {
            String[] values = params.get(key);
            if (values == null || values.length == 0) {
                continue;
            }
            for (String value : values) {
                list.add(key + "=" + URLEncoder.encode(value, "utf-8"));
            }
        }
        return CollectionUtil.join(list, "&");
    }

    /**
     * 获取Post请求的表单参数
     *
     * @param params 请求参数
     * @return
     * @throws UnsupportedEncodingException
     */
    private static UrlEncodedFormEntity getFormParam(Map<String, String[]> params) throws UnsupportedEncodingException {
        if (params == null || params.size() == 0) {
            return null;
        }
        List<NameValuePair> pairs = new ArrayList<NameValuePair>();
        for (String key : params.keySet()) {
            String[] values = params.get(key);
            if (values == null || values.length == 0) {
                continue;
            }
            for (String value : values) {
                pairs.add(new BasicNameValuePair(key, value));
            }
        }
        return new UrlEncodedFormEntity(pairs);
    }

    /**
     * 设置请求头
     *
     * @param request 请求对象
     * @param headers 请求头参数
     */
    private static void addHeaders(HttpRequestBase request, Map<String, String> headers) {
        if (headers == null || headers.size() == 0) {
            return;
        }
        for (String name : headers.keySet()) {
            request.addHeader(name, headers.get(name));
        }
    }

    /**
     * 设置cookie
     *
     * @param request 请求对象
     * @param cookies cookie参数
     * @throws UnsupportedEncodingException
     */
    private static void addCookies(HttpRequestBase request, Map<String, String> cookies) throws UnsupportedEncodingException {
        if (cookies == null || cookies.size() == 0) {
            return;
        }
        List<String> list = new ArrayList<String>();
        for (String key : cookies.keySet()) {
            String value = URLEncoder.encode(cookies.get(key), "utf-8");
            list.add(key + "=" + value);
        }
        request.addHeader("Cookie", CollectionUtil.join(list, ";"));
    }

    /**
     * 设置ContentType
     * @param request
     * @param contentType
     */
    private static void addContentType(HttpRequestBase request, String contentType) {
        request.addHeader(HttpHeaders.CONTENT_TYPE, contentType);
    }

    /**
     * 构建HttpClient对象
     *
     * @param proxy
     * @return
     */
    private static CloseableHttpClient buildHttpClient(Proxy proxy) {
        if (proxy != null && proxy.isAuthEnabled()) {
            CredentialsProvider credsProvider = new BasicCredentialsProvider();
            UsernamePasswordCredentials credentials = new UsernamePasswordCredentials(proxy.getUsername(), proxy.getPassword());
            AuthScope scope = new AuthScope(proxy.getHost(), proxy.getPort());
            credsProvider.setCredentials(scope, credentials);
            return HttpClients.custom().setDefaultCredentialsProvider(credsProvider).build();
        } else {
            return HttpClients.createDefault();
        }
    }

    /**
     * 构建RequestConfig对象
     *
     * @param proxy
     * @param timeout
     * @return
     */
    private static RequestConfig buildRequestConfig(Proxy proxy, int timeout) {
        RequestConfig.Builder builder = RequestConfig.custom();
        if (proxy != null) {
            HttpHost host = new HttpHost(proxy.getHost(), proxy.getPort(), proxy.getProtocol().getValue());
            builder.setProxy(host);
        }
        builder.setConnectionRequestTimeout(timeout).setConnectTimeout(timeout).setSocketTimeout(timeout);
        return builder.build();
    }
}
